#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author： jenny
# datetime： 2021/5/6 15:53 
# File ：stomp_ws.py
import websocket
import time
from threading import Thread
import uuid
from constants import *
BYTE = {
    'LF': '\x0A',
    'NULL': '\x00'
}

VERSIONS = '1.0,1.1'


class Stomp:
    def __init__(self, host, sockjs=False, wss=True):
        """
        Initialize STOMP communication. This is the high level API that is exposed to clients.

        Args:
            host: Hostname
            sockjs: True if the STOMP server is sockjs
            wss: True if communication is over SSL
        """
        # websocket.enableTrace(True)
        ws_host = host if sockjs is False else host + "/websocket"
        protocol = "ws://" if wss is False else "wss://"

        self.url = protocol + ws_host
        print("websocket url:"+self.url)
        self.dispatcher = Dispatcher(self)

        # maintain callback registry for subscriptions -> topic (str) vs callback (func)
        self.callback_registry = {}
        self.on_error = None
        self.on_connect = None
        self.on_message = None
        self.on_close = None

    def connect(self,username=None,passcode=None):
        """
        Connect to the remote STOMP server
        """
        # set flag to false
        self.connected = False

        # attempt to connect
        self.dispatcher.connect(username,passcode)

        # wait until connected
        start_time = time.time()
        timeout = 10  # 10 seconds
        while self.connected is False:
            if time.time() - start_time > timeout:
                print("Connection timed out")
                return False
            time.sleep(.50)
        if self.on_connect is not None:
            self.on_connect(self.connected)
            
        return self.connected

    def disconnect(self):
        """
        Disconnect from the remote STOMP server
        """        
        self.dispatcher.ws.close()
        self.connected = False
        if self.on_close is not None:
            self.on_close()
            
    def subscribe(self, destination,id=None,ack='auto',callback=None):
        """
        Subscribe to a destination and supply a callback that should be executed when a message is received on that destination
        """
        # create entry in registry against destination
        if callback is not None:
            self.callback_registry[destination] = callback

        # transmit subscribe frame
        self.dispatcher.subscribe(destination,id,ack)

    def send(self, destination, message):
        """
        Send a message to a destination
        """
        self.dispatcher.send(destination, message)


class Dispatcher:
    def __init__(self, stomp):
        """
        The Dispatcher handles all network I/O and frame marshalling/unmarshalling
        """
        self.stomp = stomp
        #websocket.enableTrace(True)  # 开启调试信息
        self.ws = websocket.WebSocketApp(self.stomp.url)

        self.ws.ping_interval = 30
        self.ws.ping_timeout = 10
        # register websocket callbacks
        self.ws.on_open = self._on_open
        self.ws.on_message = self._on_message
        self.ws.on_error = self._on_error
        self.ws.on_close = self._on_close
        self.ws.on_ping = self._on_ping

        # run event loop on separate thread
        Thread(target=self.ws.run_forever,kwargs={'ping_interval': 10, 'ping_timeout': 8}).start()

        self.opened = False

        # wait until connected
        start_time = time.time()
        timeout = 10  # 10 seconds
        while self.opened is False:
            if time.time() - start_time > timeout:
                print("WebSocket Connection timeout")
                break
            time.sleep(.50)

    def _on_message(self, ws, message):
        """
        Executed when messages is received on WS
        """
        print("<<< " + message)
        if len(message) > 0:
            command, headers, body = self._parse_message(message)
            # if connected, let Stomp know
            if command == "CONNECTED":
                self.stomp.connected = True
            # if message received, call appropriate callback
            if command == "MESSAGE":
                # 检查字典中是否存在该主题的回调函数
                if headers['destination'] in self.stomp.callback_registry:
                    self.stomp.callback_registry[headers['destination']](body)
            # if message is acked, let Stomp know
            if command == CMD_ACK:
                print("ACK: " + headers['id'])
                
            if command != '':
                if self.stomp.on_message is not None:
                    self.stomp.on_message(command, headers, body)
            

    def _on_error(self, ws, error):
        """
        Executed when WS connection errors out
        """
        print(error)
        if hasattr(self.stomp, 'on_error') and self.stomp.on_error is not None:
            self.stomp.on_error(error)

    def _on_close(self,ws,code,reason):
        """
        Executed when WS connection is closed
        """
        print("### closed ###")
        if hasattr(self.stomp, 'on_close') and self.stomp.on_close is not None:
            self.stomp.on_close(code,reason)

    def _on_open(self, ws):
        """
        Executed when WS connection is opened
        """
        self.opened = True
    
    def _on_ping(self,ws,message):
        print("### ping ###")

    def _transmit(self, command, headers, msg=None):
        """
        Marshalls and transmits the frame
        """
        # Contruct the frame
        lines = []
        lines.append(command + BYTE['LF'])

        # add headers
        for key in headers:
            lines.append(key + ":" + headers[key] + BYTE['LF'])

        lines.append(BYTE['LF'])

        # add message, if any
        if msg is not None:
            lines.append(msg)

        # terminate with null octet
        lines.append(BYTE['NULL'])

        frame = ''.join(lines)

        # transmit over ws
        print(">>>" + frame)
        if not self.dispatcher.ws.keep_running:
            print("Cannot send message: connection is closed.")
            return
        self.ws.send(frame)

    def _parse_message(self, frame):
        """
        Returns:
            command
            headers
            body

        Args:
            frame: raw frame string
        """
        lines = frame.split(BYTE['LF'])

        command = lines[0].strip()
        headers = {}

        # get all headers
        i = 1
        while lines[i] != '':
            # get key, value from raw header
            (key, value) = lines[i].split(':')
            headers[key] = value
            i += 1

        # set body to None if there is no body
        if i < len(lines) - 1:
            body = None if lines[i+1] == BYTE['NULL'] else  ''.join(lines[i+1:len(lines)-1])
            if body is not None:
                body = body.replace('\x00', '')
        else:
            body = None

        return command, headers, body

    def connect(self,username=None,passcode=None):
        """
        Transmit a CONNECT frame
        """
        headers = {}

        headers[HDR_HOST] = '/'
        headers[HDR_ACCEPT_VERSION] = VERSIONS
        headers[HDR_HEARTBEAT] = '10000,10000'
        
        if username is not None:
            headers[HDR_LOGIN] = username

        if passcode is not None:
            headers[HDR_PASSCODE] = passcode

        self._transmit(CMD_CONNECT, headers)

    def subscribe(self,destination,id,ack):
        """
        Transmit a SUBSCRIBE frame
        """
        headers = {}

        # TODO id should be auto generated
        if id is None:
            id = str(uuid.uuid4())
            
        headers[HDR_ID] = id
        headers[CMD_ACK] = ack
        headers[HDR_DESTINATION] = destination

        self._transmit(CMD_SUBSCRIBE, headers)

    def send(self, destination, message):
        """
        Transmit a SEND frame
        """
        headers = {}

        headers[HDR_DESTINATION] = destination
        headers[HDR_CONTENT_LENGTH] = str(len(message))

        self._transmit(CMD_SEND, headers, msg=message)

        
    def ack(self, message_id, subscription):
        """
        Transmit an ACK frame
        ACK 命令用于确认消息已成功处理
        当客户端接收到消息时，消息的头部会包含 message-id 字段。客户端需要从这个字段中提取 message_id
        在订阅消息时，客户端会指定一个 id，这个 id 就是 subscription
        """
        headers = {}

        headers['id'] = message_id
        headers['subscription'] = subscription

        self._transmit(CMD_ACK, headers)


def do_thing_a(msg):
    print("MESSAGE: " + msg)


def main(url,*sub_topic, **send_topic):
    stomp = Stomp(url, sockjs=False, wss=True)
    stomp.connect()
    stomp.subscribe(sub_topic, do_thing_a)
    time.sleep(2)
    stomp.send(send_topic, '{"name":"akshaye"}')
if __name__ == "__main__":
	main()
